"""
特征提取器主类

整合时域和频域特征提取功能
"""

import numpy as np
from typing import Dict, List, Optional
from loguru import logger
from .time_domain_features import TimeDomainFeatures
from .frequency_domain_features import FrequencyDomainFeatures


class FeatureExtractor:
    """特征提取器主类"""
    
    def __init__(self, 
                 time_features: Optional[List[str]] = None,
                 freq_features: Optional[List[str]] = None,
                 sample_rate: int = 100):
        """
        初始化特征提取器
        
        Args:
            time_features: 时域特征列表
            freq_features: 频域特征列表
            sample_rate: 采样率
        """
        self.time_extractor = TimeDomainFeatures(time_features)
        self.freq_extractor = FrequencyDomainFeatures(freq_features, sample_rate)
        self.sample_rate = sample_rate
        
        logger.info("特征提取器初始化完成")
    
    def extract_all_features(self, signal_data: np.ndarray) -> Dict[str, np.ndarray]:
        """
        提取所有特征
        
        Args:
            signal_data: 输入信号
            
        Returns:
            所有特征字典
        """
        # 提取时域特征
        time_features = self.time_extractor.extract_features(signal_data)
        
        # 提取频域特征
        freq_features = self.freq_extractor.extract_features(signal_data)
        
        # 合并特征
        all_features = {**time_features, **freq_features}
        
        return all_features
    
    def extract_feature_vector(self, signal_data: np.ndarray) -> np.ndarray:
        """
        提取特征向量
        
        Args:
            signal_data: 输入信号
            
        Returns:
            特征向量
        """
        # 提取所有特征
        features = self.extract_all_features(signal_data)
        
        # 转换为向量
        feature_vector = []
        for feature_name, feature_value in features.items():
            if isinstance(feature_value, np.ndarray):
                # 对于数组特征，取均值或前几个值
                if feature_name == "power_spectral_density":
                    feature_vector.extend(feature_value[:20])  # 取前20个值
                elif feature_name == "mfcc":
                    feature_vector.extend(feature_value)  # 取所有MFCC系数
                else:
                    feature_vector.append(np.mean(feature_value))
            else:
                feature_vector.append(feature_value)
        
        return np.array(feature_vector)
    
    def get_feature_names(self) -> List[str]:
        """
        获取特征名称列表
        
        Returns:
            特征名称列表
        """
        time_names = self.time_extractor.get_feature_names()
        freq_names = self.freq_extractor.get_feature_names()
        
        # 处理频域特征名称（数组特征需要展开）
        expanded_freq_names = []
        for name in freq_names:
            if name == "power_spectral_density":
                expanded_freq_names.extend([f"psd_{i}" for i in range(20)])
            elif name == "mfcc":
                expanded_freq_names.extend([f"mfcc_{i}" for i in range(13)])  # 假设13个MFCC系数
            else:
                expanded_freq_names.append(name)
        
        return time_names + expanded_freq_names
    
    def get_feature_dimensions(self) -> Dict[str, int]:
        """
        获取特征维度
        
        Returns:
            特征维度字典
        """
        time_dims = {name: 1 for name in self.time_extractor.get_feature_names()}
        freq_dims = self.freq_extractor.get_feature_dimensions()
        
        # 合并维度信息
        all_dims = {**time_dims, **freq_dims}
        
        return all_dims
    
    def __str__(self) -> str:
        return f"FeatureExtractor(sample_rate={self.sample_rate})"
    
    def __repr__(self) -> str:
        return self.__str__() 